#include "USBPort.h"

#include <setupapi.h>

#define BUFFER_SIZE 2048

USBPort::USBPort(void)
: m_pUSBDevice(NULL), m_pUSBHub(NULL)
{
}

USBPort::~USBPort(void)
{
	DeleteDevice();
	DeleteHub();
}

USBDevice* USBPort::GetDevice()
{
	DeleteDevice();

	if (m_portIsDeviceConnected)
	{
		m_pUSBDevice = new USBDevice();
		if (m_pUSBDevice)
		{
			m_pUSBDevice->SetPortNumber(GetPortNumber());
			m_pUSBDevice->SetHubDevicePath(GetHubDevicePath().c_str());
			m_pUSBDevice->SetDeviceDescriptor(*GetDeviceDescriptor());

			HANDLE hHub = CreateFile(GetHubDevicePath().c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
			if (hHub != INVALID_HANDLE_VALUE)
			{
				BYTE *buffer = new BYTE[BUFFER_SIZE];
				if (buffer)
				{
					if (GetDeviceDescriptor()->iManufacturer > 0)
					{
						m_pUSBDevice->SetManufacturer(GetDescriptorString(hHub, GetDeviceDescriptor()->iManufacturer).c_str());
					}

					if (GetDeviceDescriptor()->iProduct > 0)
					{
						m_pUSBDevice->SetProduct(GetDescriptorString(hHub, GetDeviceDescriptor()->iProduct).c_str());
					}

					if (GetDeviceDescriptor()->iSerialNumber > 0)
					{
						m_pUSBDevice->SetSerialNumber(GetDescriptorString(hHub, GetDeviceDescriptor()->iSerialNumber).c_str());
					}

					memset(buffer, 0, BUFFER_SIZE);
					USB_NODE_CONNECTION_DRIVERKEY_NAME *pDriverKey = (USB_NODE_CONNECTION_DRIVERKEY_NAME*)buffer;
					pDriverKey->ConnectionIndex = GetPortNumber();
					DWORD bytesReturned;
					if (DeviceIoControl(hHub, IOCTL_USB_GET_NODE_CONNECTION_DRIVERKEY_NAME, pDriverKey, BUFFER_SIZE, pDriverKey, BUFFER_SIZE, &bytesReturned, NULL))
					{
						m_pUSBDevice->SetDriverKey(pDriverKey->DriverKeyName);
						m_pUSBDevice->SetName(GetDescriptionByKeyName(pDriverKey->DriverKeyName).c_str());
						m_pUSBDevice->SetInstanceID(GetInstanceIDByKeyName(pDriverKey->DriverKeyName).c_str());
					}

					delete [] buffer;
				}

				CloseHandle(hHub);
			}
		}
	}

	return m_pUSBDevice;
}

USBHub* USBPort::GetHub()
{
	DeleteHub();

	if (m_portIsHub)
	{
		m_pUSBHub = new USBHub();
		if (m_pUSBHub)
		{
			m_pUSBHub->SetRootHub(false);
			m_pUSBHub->SetDeviceDescription(L"External Hub");

			HANDLE hHub = CreateFile(GetHubDevicePath().c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
			if (hHub != INVALID_HANDLE_VALUE)
			{
				BYTE *pBuffer = new BYTE[BUFFER_SIZE];
				if (pBuffer)
				{
					memset(pBuffer, 0, BUFFER_SIZE);
					USB_NODE_CONNECTION_NAME *pNodeName = (USB_NODE_CONNECTION_NAME*)pBuffer;
					pNodeName->ConnectionIndex = GetPortNumber();

					DWORD bytesReturned;
					if (DeviceIoControl(hHub, IOCTL_USB_GET_NODE_CONNECTION_NAME, pNodeName, BUFFER_SIZE, pNodeName, BUFFER_SIZE, &bytesReturned, NULL))
					{
						std::wstring hubDevicePath = L"\\\\.\\";
						hubDevicePath += pNodeName->NodeName;
						m_pUSBHub->SetDevicePath(hubDevicePath.c_str());

						HANDLE hSubHub = CreateFile(m_pUSBHub->GetDevicePath().c_str(), GENERIC_WRITE, FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);
						if (hSubHub != INVALID_HANDLE_VALUE)
						{
							USB_NODE_INFORMATION nodeInformation;
							nodeInformation.NodeType = UsbHub;

							if (DeviceIoControl(hSubHub, IOCTL_USB_GET_NODE_INFORMATION, &nodeInformation, sizeof(nodeInformation), &nodeInformation, sizeof(nodeInformation), &bytesReturned, NULL))
							{
								m_pUSBHub->SetBusPowered(nodeInformation.u.HubInformation.HubIsBusPowered ? true : false);
								m_pUSBHub->SetPortCount(nodeInformation.u.HubInformation.HubDescriptor.bNumberOfPorts);
							}

							CloseHandle(hSubHub);
						}
					}

					delete [] pBuffer;
				}

				CloseHandle(hHub);
			}

			USBDevice *pUSBDevice = GetDevice();
			m_pUSBHub->SetInstanceID(pUSBDevice->GetInstanceID().c_str());
			m_pUSBHub->SetManufacturer(pUSBDevice->GetManufacturer().c_str());
			m_pUSBHub->SetProduct(pUSBDevice->GetProduct().c_str());
			m_pUSBHub->SetSerialNumber(pUSBDevice->GetSerialNumber().c_str());
			m_pUSBHub->SetDriverKey(pUSBDevice->GetDriverKey().c_str());
		}
	}

	return m_pUSBHub;
}

void USBPort::DeleteDevice()
{
	if (m_pUSBDevice)
	{
		delete m_pUSBDevice;
		m_pUSBDevice = NULL;
	}
}

void USBPort::DeleteHub()
{
	if (m_pUSBHub)
	{
		delete m_pUSBHub;
		m_pUSBHub = NULL;
	}
}

std::wstring USBPort::GetDescriptorString(HANDLE hHub, UCHAR descriptorCode)
{
	std::wstring descriptorString;
	BYTE *pBuffer = new BYTE[BUFFER_SIZE];
	if (pBuffer)
	{
		memset(pBuffer, 0, BUFFER_SIZE);

		USB_DESCRIPTOR_REQUEST *pRequest = (USB_DESCRIPTOR_REQUEST*)pBuffer;
		pRequest->ConnectionIndex = GetPortNumber();
		pRequest->SetupPacket.wValue = (USHORT)((USB_STRING_DESCRIPTOR_TYPE << 8) + descriptorCode);
		pRequest->SetupPacket.wLength = (USHORT)(BUFFER_SIZE - sizeof(USB_DESCRIPTOR_REQUEST));
		pRequest->SetupPacket.wIndex = 0x0409; // Language Code

		DWORD bytesReturned;
		if (DeviceIoControl(hHub, IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION, pRequest, BUFFER_SIZE, pRequest, BUFFER_SIZE, &bytesReturned, NULL))
		{
			USB_STRING_DESCRIPTOR *pUsbStringDescriptor = (USB_STRING_DESCRIPTOR*)pRequest->Data;
			descriptorString = pUsbStringDescriptor->bString;
		}

		delete [] pBuffer;
	}

	return descriptorString;
}

std::wstring USBPort::GetDescriptionByKeyName(wchar_t *driverKeyName)
{
	const wchar_t *REGSTR_KEY_USB = L"USB";
	std::wstring description;

	HDEVINFO hDevInfo = SetupDiGetClassDevs(NULL, REGSTR_KEY_USB, NULL, DIGCF_PRESENT | DIGCF_ALLCLASSES);
	if (hDevInfo != INVALID_HANDLE_VALUE)
	{
		BYTE *pBuffer = new BYTE[BUFFER_SIZE];
		if (pBuffer)
		{
			BOOL success;
			int i = 0;

			do 
			{
				SP_DEVINFO_DATA sdd;
				sdd.cbSize = sizeof(sdd);

				success = SetupDiEnumDeviceInfo(hDevInfo, i, &sdd);
				if (success)
				{
					DWORD requiredSize;
					DWORD regType;

					if (SetupDiGetDeviceRegistryProperty(hDevInfo, &sdd, SPDRP_DRIVER, &regType, pBuffer, BUFFER_SIZE, &requiredSize))
					{
						wchar_t *pKeyName = (wchar_t*)pBuffer;
						if (wcscmp(pKeyName, driverKeyName) == 0)
						{
							if (SetupDiGetDeviceRegistryProperty(hDevInfo, &sdd, SPDRP_DEVICEDESC, &regType, pBuffer, BUFFER_SIZE, &requiredSize))
							{
								description = (wchar_t*)pBuffer;
							}
							break;
						}
					}
				}
				++i;
			} while (success);

			delete [] pBuffer;
		}

		SetupDiDestroyDeviceInfoList(hDevInfo);
	}

	return description;
}

std::wstring USBPort::GetInstanceIDByKeyName(wchar_t *driverKeyName)
{
	const wchar_t *REGSTR_KEY_USB = L"USB";
	std::wstring instanceID;

	HDEVINFO hDevInfo = SetupDiGetClassDevs(NULL, REGSTR_KEY_USB, NULL, DIGCF_PRESENT | DIGCF_ALLCLASSES);
	if (hDevInfo != INVALID_HANDLE_VALUE)
	{
		BYTE *pBuffer = new BYTE[BUFFER_SIZE];
		if (pBuffer)
		{
			BOOL success;
			int i = 0;

			do 
			{
				SP_DEVINFO_DATA sdd;
				sdd.cbSize = sizeof(sdd);

				success = SetupDiEnumDeviceInfo(hDevInfo, i, &sdd);
				if (success)
				{
					DWORD requiredSize;
					DWORD regType;

					if (SetupDiGetDeviceRegistryProperty(hDevInfo, &sdd, SPDRP_DRIVER, &regType, pBuffer, BUFFER_SIZE, &requiredSize))
					{
						wchar_t *pKeyName = (wchar_t*)pBuffer;
						if (wcscmp(pKeyName, driverKeyName) == 0)
						{
							if (SetupDiGetDeviceInstanceId(hDevInfo, &sdd, (wchar_t*)pBuffer, BUFFER_SIZE, &requiredSize))
							{
								instanceID = (wchar_t*)pBuffer;
							}
							break;
						}
					}
				}
				++i;
			} while (success);

			delete [] pBuffer;
		}

		SetupDiDestroyDeviceInfoList(hDevInfo);
	}

	return instanceID;
}